#include "common.h"

// 客户端处理线程函数
void *handle_client(void *client_socket_ptr) {
    int client_socket = *(int *)client_socket_ptr;
    free(client_socket_ptr);

    printf("New client connected\n");

    while (1) {
        MessageHeader header;
        void *data = NULL;

        // 接收客户端消息
        ssize_t recv_len = receive_message(client_socket, &header, &data);
        if (recv_len <= 0) {
            if (recv_len < 0)
                perror("receive_message");
            break;
        }

        switch (header.type) {
            case MSG_FILE_REQUEST: {
                // 客户端请求文件信息
                char *filename = (char *)data;
                printf("Client requested file: %s\n", filename);

                // 检查文件是否存在
                struct stat file_stat;
                if (stat(filename, &file_stat) == -1) {
                    // 文件不存在，发送错误消息
                    TransferStatusMsg status_msg;
                    strncpy(status_msg.filename, filename, MAX_FILENAME_LENGTH);
                    status_msg.status = STATUS_ERROR;
                    snprintf(status_msg.message, sizeof(status_msg.message),
                             "File not found: %s", filename);
                    send_message(client_socket, MSG_TRANSFER_STATUS, &status_msg, sizeof(status_msg));
                    free(data);
                    break;
                }

                // 计算文件信息
                FileInfo file_info;
                strncpy(file_info.filename, filename, MAX_FILENAME_LENGTH);
                file_info.file_size = file_stat.st_size;
                file_info.block_size = DEFAULT_BLOCK_SIZE;
                file_info.total_blocks = (file_info.file_size + file_info.block_size - 1) / file_info.block_size;

                // 计算文件MD5
                if (calculate_file_md5(filename, file_info.md5) == -1) {
                    // MD5计算失败，发送错误消息
                    TransferStatusMsg status_msg;
                    strncpy(status_msg.filename, filename, MAX_FILENAME_LENGTH);
                    status_msg.status = STATUS_ERROR;
                    snprintf(status_msg.message, sizeof(status_msg.message),
                             "Failed to calculate MD5 for file: %s", filename);
                    send_message(client_socket, MSG_TRANSFER_STATUS, &status_msg, sizeof(status_msg));
                    free(data);
                    break;
                }

                // 发送文件信息
                send_message(client_socket, MSG_FILE_INFO, &file_info, sizeof(FileInfo));
                free(data);
                break;
            }

            case MSG_BLOCK_REQUEST: {
                // 客户端请求块数据
                BlockRequest *request = (BlockRequest *)data;
                printf("Client requested block %d of file: %s\n", request->block_index, request->filename);

                // 打开文件
                FILE *file = fopen(request->filename, "rb");
                if (!file) {
                    // 文件打开失败，发送错误消息
                    TransferStatusMsg status_msg;
                    strncpy(status_msg.filename, request->filename, MAX_FILENAME_LENGTH);
                    status_msg.status = STATUS_ERROR;
                    snprintf(status_msg.message, sizeof(status_msg.message),
                             "Failed to open file: %s", request->filename);
                    send_message(client_socket, MSG_TRANSFER_STATUS, &status_msg, sizeof(status_msg));
                    free(data);
                    break;
                }

                // 获取文件大小
                struct stat file_stat;
                if (fstat(fileno(file), &file_stat) == -1) {
                    // 获取文件大小失败，发送错误消息
                    TransferStatusMsg status_msg;
                    strncpy(status_msg.filename, request->filename, MAX_FILENAME_LENGTH);
                    status_msg.status = STATUS_ERROR;
                    snprintf(status_msg.message, sizeof(status_msg.message),
                             "Failed to get file size: %s", request->filename);
                    send_message(client_socket, MSG_TRANSFER_STATUS, &status_msg, sizeof(status_msg));
                    fclose(file);
                    free(data);
                    break;
                }

                // 计算块信息
                int block_size = DEFAULT_BLOCK_SIZE;
                int total_blocks = (file_stat.st_size + block_size - 1) / block_size;

                if (request->block_index < 0 || request->block_index >= total_blocks) {
                    // 块索引无效，发送错误消息
                    TransferStatusMsg status_msg;
                    strncpy(status_msg.filename, request->filename, MAX_FILENAME_LENGTH);
                    status_msg.status = STATUS_ERROR;
                    snprintf(status_msg.message, sizeof(status_msg.message),
                             "Invalid block index %d (total blocks: %d)", request->block_index, total_blocks);
                    send_message(client_socket, MSG_TRANSFER_STATUS, &status_msg, sizeof(status_msg));
                    fclose(file);
                    free(data);
                    break;
                }

                // 计算块偏移和大小
                off_t offset = (off_t)request->block_index * block_size;
                size_t bytes_to_read = block_size;
                if (offset + bytes_to_read > file_stat.st_size) {
                    bytes_to_read = file_stat.st_size - offset;
                }

                // 定位到块位置
                if (fseek(file, offset, SEEK_SET) != 0) {
                    // 定位失败，发送错误消息
                    TransferStatusMsg status_msg;
                    strncpy(status_msg.filename, request->filename, MAX_FILENAME_LENGTH);
                    status_msg.status = STATUS_ERROR;
                    snprintf(status_msg.message, sizeof(status_msg.message),
                             "Failed to seek to block %d", request->block_index);
                    send_message(client_socket, MSG_TRANSFER_STATUS, &status_msg, sizeof(status_msg));
                    fclose(file);
                    free(data);
                    break;
                }

                // 分配块数据缓冲区
                BlockData *block_data = malloc(sizeof(BlockData) + bytes_to_read);
                if (!block_data) {
                    // 内存分配失败，发送错误消息
                    TransferStatusMsg status_msg;
                    strncpy(status_msg.filename, request->filename, MAX_FILENAME_LENGTH);
                    status_msg.status = STATUS_ERROR;
                    snprintf(status_msg.message, sizeof(status_msg.message),
                             "Memory allocation failed for block data");
                    send_message(client_socket, MSG_TRANSFER_STATUS, &status_msg, sizeof(status_msg));
                    fclose(file);
                    free(data);
                    break;
                }

                // 读取块数据
                size_t bytes_read = fread(block_data->data, 1, bytes_to_read, file);
                if (bytes_read != bytes_to_read) {
                    // 读取失败，发送错误消息
                    TransferStatusMsg status_msg;
                    strncpy(status_msg.filename, request->filename, MAX_FILENAME_LENGTH);
                    status_msg.status = STATUS_ERROR;
                    snprintf(status_msg.message, sizeof(status_msg.message),
                             "Failed to read block %d (expected %zu bytes, got %zu)",
                             request->block_index, bytes_to_read, bytes_read);
                    send_message(client_socket, MSG_TRANSFER_STATUS, &status_msg, sizeof(status_msg));
                    free(block_data);
                    fclose(file);
                    free(data);
                    break;
                }

                // 填充块数据信息
                strncpy(block_data->filename, request->filename, MAX_FILENAME_LENGTH);
                block_data->block_index = request->block_index;
                block_data->block_size = bytes_read;

                // 计算块MD5
                calculate_block_md5(block_data->data, bytes_read, block_data->md5);

                // 发送块数据
                send_message(client_socket, MSG_BLOCK_DATA, block_data, sizeof(BlockData) + bytes_read);

                // 清理
                free(block_data);
                fclose(file);
                free(data);
                break;
            }

            default:
                printf("Unknown message type: %d\n", header.type);
                free(data);
                break;
        }
    }

    printf("Client disconnected\n");
    close(client_socket);
    return NULL;
}

int main(int argc, char *argv[]) {
    int port = DEFAULT_PORT;
    if (argc > 1) {
        port = atoi(argv[1]);
    }

    // 创建监听套接字
    int listen_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (listen_socket < 0) {
        error_exit("socket");
    }

    // 设置套接字选项
    int opt = 1;
    if (setsockopt(listen_socket, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
        error_exit("setsockopt");
    }

    // 绑定地址和端口
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(port);

    if (bind(listen_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        error_exit("bind");
    }

    // 开始监听
    if (listen(listen_socket, MAX_CLIENTS) < 0) {
        error_exit("listen");
    }

    printf("Server listening on port %d...\n", port);

    while (1) {
        // 接受客户端连接
        struct sockaddr_in client_addr;
        socklen_t client_addr_len = sizeof(client_addr);
        int *client_socket = malloc(sizeof(int));
        *client_socket = accept(listen_socket, (struct sockaddr *)&client_addr, &client_addr_len);
        if (*client_socket < 0) {
            perror("accept");
            free(client_socket);
            continue;
        }

        // 创建线程处理客户端
        pthread_t thread_id;
        if (pthread_create(&thread_id, NULL, handle_client, client_socket) != 0) {
            perror("pthread_create");
            close(*client_socket);
            free(client_socket);
            continue;
        }

        // 分离线程，自动释放资源
        pthread_detach(thread_id);
    }

    // 关闭监听套接字
    close(listen_socket);
    return 0;
}
